from flask_sqlalchemy import SQLAlchemy, SignallingSession, get_state
from sqlalchemy import orm

class RoutingSession(SignallingSession):
    def __init__(self,*args,**kwargs):
        super(RoutingSession,self).__init__(*args,**kwargs)

    def get_bind(self, mapper=None, clause=None):
        """根据数据库操作定位数据库引擎"""
        state = get_state(self.app)
        if mapper is not None:
            try:
                # SA >= 1.3
                persist_selectable = mapper.persist_selectable
            except AttributeError:
                # SA < 1.3
                persist_selectable = mapper.mapped_table

            info = getattr(persist_selectable, 'info', {})
            bind_key = info.get('bind_key')
            if bind_key is not None:
                state = get_state(self.app)
                return state.db.get_engine(self.app, bind=bind_key)
        from sqlalchemy.sql.dml import Update,Insert,Delete
        if self._flushing or isinstance(clause,(Update,Insert,Delete)):
            print("读取主库")
            print(self._flushing)
            return state.db.get_engine(self.app,bind='master')
        else:
            print("读取从库")
            return state.db.get_engine(self.app, bind='slave1')

class RoutingSQlAlchemy(SQLAlchemy):

    def create_session(self, options):
        return orm.sessionmaker(class_=RoutingSession, db=self, **options)
